from ultralytics import YOLO


class Detector:
    def __init__(self, model_path, conf_thres=0.5):
        """初始化目标检测器"""
        self.model = YOLO(model_path)
        self.conf_thres = conf_thres

    def detect(self, image):
        """执行目标检测"""
        results = self.model(image, conf=self.conf_thres)
        detections = []

        for r in results:
            boxes = r.boxes
            for box in boxes:
                x1, y1, x2, y2 = box.xyxy[0]
                conf = box.conf[0]
                cls = box.cls[0]

                detections.append(
                    {
                        "bbox": [x1, y1, x2, y2],
                        "confidence": conf,
                        "class_id": cls,
                    }
                )

        return detections
